import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from PIL import Image
import os
import matplotlib.pyplot as plt
import numpy as np
import random


class Config:
    IMAGE_SIZE = (128, 128)
    BATCH_SIZE = 64
    NUM_EPOCHS = 5
    LEARNING_RATE = 0.001
    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    MAX_SAMPLES_PER_CLASS = 500
    DATA_DIR = r'PetImages'  # 使用原始字符串确保路径正确


class PetDataset(Dataset):
    def __init__(self, root_dir, transform=None, max_samples_per_class=None):
        self.root_dir = root_dir
        self.transform = transform
        self.classes = ['Cat', 'Dog']
        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}

        self.images = []
        self.labels = []

        # 调试：打印目录信息
        print(f"数据集根目录: {os.path.abspath(root_dir)}")
        print(f"目录存在: {os.path.exists(root_dir)}")

        for class_name in self.classes:
            class_dir = os.path.join(root_dir, class_name)
            print(f"检查类别目录: {class_dir}")
            print(f"目录存在: {os.path.exists(class_dir)}")

            if not os.path.exists(class_dir):
                print(f"警告: {class_dir} 不存在!")
                continue

            img_files = [f for f in os.listdir(class_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
            print(f"{class_name} 找到 {len(img_files)} 张图片")

            if max_samples_per_class and len(img_files) > max_samples_per_class:
                img_files = random.sample(img_files, max_samples_per_class)

            for img_name in img_files:
                img_path = os.path.join(class_dir, img_name)
                try:
                    # 验证图像可打开
                    with Image.open(img_path) as img:
                        img.verify()
                    self.images.append(img_path)
                    self.labels.append(self.class_to_idx[class_name])
                except Exception as e:
                    print(f"跳过损坏图像: {img_path} - {str(e)}")
                    continue

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_path = self.images[idx]
        label = self.labels[idx]

        try:
            image = Image.open(img_path).convert('RGB')
            if self.transform:
                image = self.transform(image)
            return image, label
        except Exception as e:
            # 返回一个默认图像
            print(f"加载错误: {img_path}, 使用灰度图替代 - {str(e)}")
            dummy_image = Image.new('RGB', Config.IMAGE_SIZE, color='gray')
            if self.transform:
                dummy_image = self.transform(dummy_image)
            return dummy_image, label


def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs):
    train_losses = []
    train_accs = []
    val_losses = []
    val_accs = []

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0

        # 简化进度条
        for images, labels in train_loader:
            images, labels = images.to(Config.DEVICE), labels.to(Config.DEVICE)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        epoch_loss = running_loss / len(train_loader)
        epoch_acc = 100 * correct / total
        train_losses.append(epoch_loss)
        train_accs.append(epoch_acc)

        # 验证
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0

        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(Config.DEVICE), labels.to(Config.DEVICE)
                outputs = model(images)
                loss = criterion(outputs, labels)

                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()

        val_epoch_loss = val_loss / len(val_loader)
        val_epoch_acc = 100 * val_correct / val_total
        val_losses.append(val_epoch_loss)
        val_accs.append(val_epoch_acc)

        print(f'Epoch {epoch + 1}/{num_epochs}:')
        print(f'Train Loss: {epoch_loss:.4f}, Train Acc: {epoch_acc:.2f}%')
        print(f'Val Loss: {val_epoch_loss:.4f}, Val Acc: {val_epoch_acc:.2f}%')

    return train_losses, train_accs, val_losses, val_accs


def plot_metrics(train_losses, train_accs, val_losses, val_accs):
    plt.figure(figsize=(10, 4))

    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Train Loss')
    plt.plot(val_losses, label='Validation Loss')
    plt.title('Loss Curves')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(train_accs, label='Train Accuracy')
    plt.plot(val_accs, label='Validation Accuracy')
    plt.title('Accuracy Curves')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.legend()

    plt.tight_layout()
    plt.savefig('training_metrics.png', dpi=100)
    plt.close()


def visualize_predictions(model, val_loader, num_samples=6):
    model.eval()
    images, labels = next(iter(val_loader))
    images = images[:num_samples].to(Config.DEVICE)
    labels = labels[:num_samples]

    with torch.no_grad():
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)

    plt.figure(figsize=(12, 4))
    for i in range(num_samples):
        plt.subplot(2, 3, i + 1)
        img = images[i].cpu().numpy().transpose((1, 2, 0))
        img = img * np.array([0.229, 0.224, 0.225]) + np.array([0.485, 0.456, 0.406])
        img = np.clip(img, 0, 1)

        plt.imshow(img)
        plt.title(f'True: {"Cat" if labels[i] == 0 else "Dog"}\nPred: {"Cat" if predicted[i] == 0 else "Dog"}')
        plt.axis('off')

    plt.tight_layout()
    plt.savefig('prediction_visualization.png', dpi=100)
    plt.close()


def main():
    # 设置随机种子
    random.seed(42)
    torch.manual_seed(42)

    # 检查数据集路径
    print(f"当前工作目录: {os.getcwd()}")
    print(f"数据集路径: {Config.DATA_DIR}")
    print(f"路径存在: {os.path.exists(Config.DATA_DIR)}")

    # 数据变换
    transform_train = transforms.Compose([
        transforms.Resize(Config.IMAGE_SIZE),
        transforms.RandomHorizontalFlip(p=0.3),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    transform_valid = transforms.Compose([
        transforms.Resize(Config.IMAGE_SIZE),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # 创建数据集
    train_dataset = PetDataset(Config.DATA_DIR, transform=transform_train,
                               max_samples_per_class=Config.MAX_SAMPLES_PER_CLASS)
    val_dataset = PetDataset(Config.DATA_DIR, transform=transform_valid,
                             max_samples_per_class=Config.MAX_SAMPLES_PER_CLASS // 4)

    print(f"训练集大小: {len(train_dataset)}")
    print(f"验证集大小: {len(val_dataset)}")

    # 如果数据集为空则退出
    if len(train_dataset) == 0 or len(val_dataset) == 0:
        print("错误: 数据集为空，请检查路径和文件格式")
        return

    # ... (其余代码保持不变)
    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=Config.BATCH_SIZE,
                              shuffle=True, num_workers=2, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=Config.BATCH_SIZE,
                            shuffle=False, num_workers=2, pin_memory=True)

    # 使用更轻量的模型
    model = models.resnet18(pretrained=True)  # 使用ResNet18而不是ResNet50
    model.fc = nn.Linear(model.fc.in_features, 2)
    model = model.to(Config.DEVICE)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=Config.LEARNING_RATE)

    print("开始训练...")
    # 训练模型
    train_losses, train_accs, val_losses, val_accs = train_model(
        model, train_loader, val_loader, criterion, optimizer, Config.NUM_EPOCHS
    )

    # 绘制训练指标
    plot_metrics(train_losses, train_accs, val_losses, val_accs)

    # 可视化预测结果
    visualize_predictions(model, val_loader)

    # 保存模型
    torch.save(model.state_dict(), 'cat_dog_classifier.pth')
    print("训练完成！模型已保存为 cat_dog_classifier.pth")


if __name__ == '__main__':
    main()